using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using System.Threading.Tasks;
using DotNetty.Buffers;
using DotNetty.Codecs;
using DotNetty.Transport.Bootstrapping;
using DotNetty.Transport.Channels;
using DotNetty.Transport.Channels.Sockets;
using Microsoft.Extensions.Logging;
using ProtoBuf;

namespace Shielded.Gossip
{
    public class NettyTransport : ITransport
    {
        public string OwnId { get; }
        public ICollection<string> Servers => ServerIPs.Keys;
        public MessageHandler MessageHandler { get; set; }
        public void Start()
        {
            StartListening();
        }

        public readonly ShieldedDict<string, IPEndPoint> ServerIPs;
        public IPEndPoint LocalEndpoint { get; }
        public int ReceiveTimeout { get; set; } = 30000;
        public int KeepAliveInterval { get; set; } = 10000;
        public event EventHandler<Exception> Error;

        private readonly ILogger _logger;
        private IChannel _serverChannel;
        private readonly ConcurrentDictionary<string, IChannel> _clientChannels = new();
        private IEventLoopGroup _bossGroup;
        private IEventLoopGroup _workerGroup;
        private readonly object _clientLock = new();
        private bool _listening = false;

        public NettyTransport(string ownId, IDictionary<string, IPEndPoint> serverIPs, ILogger logger = null)
        {
            OwnId = ownId;
            _logger = logger;
            ServerIPs = new ShieldedDict<string, IPEndPoint>(serverIPs);
            LocalEndpoint = serverIPs[ownId];
        }

        public void StartListening()
        {
            if (_listening) return;
            _listening = true;
            _bossGroup = new MultithreadEventLoopGroup(1);
            _workerGroup = new MultithreadEventLoopGroup();
            var bootstrap = new ServerBootstrap();
            bootstrap
                .Group(_bossGroup, _workerGroup)
                .Channel<TcpServerSocketChannel>()
                .Option(ChannelOption.SoBacklog, 100)
                .ChildHandler(new ActionChannelInitializer<IChannel>(channel =>
                {
                    channel.Pipeline.AddLast(new LengthFieldBasedFrameDecoder(1024 * 1024, 0, 4, 0, 4));
                    channel.Pipeline.AddLast(new LengthFieldPrepender(4));
                    channel.Pipeline.AddLast(new NettyServerHandler(this));
                }));
            _serverChannel = bootstrap.BindAsync(LocalEndpoint).Result;
        }

        public void StopListening()
        {
            _serverChannel?.CloseAsync().Wait();
            _serverChannel = null;
            _bossGroup?.ShutdownGracefullyAsync().Wait();
            _workerGroup?.ShutdownGracefullyAsync().Wait();
            _listening = false;
        }

        public void Broadcast(object msg)
        {
            foreach (var server in Servers)
            {
                if (server != OwnId)
                    Send(server, msg, false);
            }
        }

        public void Send(string server, object msg, bool replyExpected)
        {
            if (!ServerIPs.TryGetValue(server, out var endpoint))
                return;
            var channel = _clientChannels.GetOrAdd(server, _ => ConnectToPeer(endpoint));
            var bytes = Serializer.Serialize(msg);
            var buffer = Unpooled.WrappedBuffer(bytes);
            channel.WriteAndFlushAsync(buffer);
        }

        private IChannel ConnectToPeer(IPEndPoint endpoint)
        {
            lock (_clientLock)
            {
                var group = new MultithreadEventLoopGroup();
                var bootstrap = new Bootstrap();
                bootstrap
                    .Group(group)
                    .Channel<TcpSocketChannel>()
                    .Option(ChannelOption.TcpNodelay, true)
                    .Handler(new ActionChannelInitializer<IChannel>(channel =>
                    {
                        channel.Pipeline.AddLast(new LengthFieldBasedFrameDecoder(1024 * 1024, 0, 4, 0, 4));
                        channel.Pipeline.AddLast(new LengthFieldPrepender(4));
                        channel.Pipeline.AddLast(new NettyClientHandler(this));
                    }));
                return bootstrap.ConnectAsync(endpoint).Result;
            }
        }

        public void Dispose()
        {
            StopListening();
            foreach (var ch in _clientChannels.Values)
                ch.CloseAsync().Wait();
        }

        public void RaiseError(Exception ex) => Error?.Invoke(this, ex);

        // Handler for server-side
        class NettyServerHandler : SimpleChannelInboundHandler<IByteBuffer>
        {
            private readonly NettyTransport _transport;
            public NettyServerHandler(NettyTransport transport) => _transport = transport;
            protected override void ChannelRead0(IChannelHandlerContext ctx, IByteBuffer msg)
            {
                try
                {
                    var bytes = new byte[msg.ReadableBytes];
                    msg.ReadBytes(bytes);
                    var obj = Serializer.Deserialize(bytes); 
                    var reply = _transport.MessageHandler?.Invoke(obj);
                    // if (reply != null)
                    // {
                    //     var replyBytes = Serializer.Serialize(reply);
                    //     var buffer = Unpooled.WrappedBuffer(replyBytes);
                    //     ctx.WriteAndFlushAsync(buffer);
                    // }
                }
                catch (Exception ex)
                {
                    _transport.RaiseError(ex);
                }
            }
        }

        // Handler for client-side (optional, for reply handling)
        class NettyClientHandler : SimpleChannelInboundHandler<IByteBuffer>
        {
            private readonly NettyTransport _transport;
            public NettyClientHandler(NettyTransport transport) => _transport = transport;
            protected override void ChannelRead0(IChannelHandlerContext ctx, IByteBuffer msg)
            {
                // 可选：处理reply
            }
            public override void ExceptionCaught(IChannelHandlerContext context, Exception exception)
            {
                _transport.RaiseError(exception);
            }
        }
    }

    [ProtoContract]
    public class BrokerMessage
    {
        [ProtoMember(1)]
        public int flag;
        [ProtoMember(2)]
        public byte[] bytes;
    }
} 